import { Academy } from '../Academy';
import { ActionSpec } from '../Actuators/ActionSpec';
import { ActionBuffers } from '../Actuators/IActionReceiver';
import { AgentInfo } from '../Agent';
import {
	RewardSignals,
	TrainingBehaviorInitializedEvent,
	TrainingEnvironmentInitializedEvent,
	TrainingFeatures,
} from '../Analytics/Events';
import { DemonstrationMetaData } from '../Demonstrations/DemonstrationMetaData';
import pb from '../Grpc/CommunicatorObjects/pb.js';
import { BrainParameters } from '../Policies/BrainParameters';
import { SensorCompressionType } from '../Sensors/CompressionSpec';
import { DimensionProperty, ISensor, SensorExtensions } from '../Sensors/ISensor';
import { ObservationWriter } from '../Sensors/ObservationWriter';
import { CCCRLCapabilities } from './CCCRLCapabilities';
import { CCCRLInitParameters } from './ICommunicator';

export class GrpcExtensions {
	private static hasWarnedAboutMultiAgentGroups = false;

	static capabilitiesToProto(rlCaps: CCCRLCapabilities): pb.communicator_objects.UnityRLCapabilitiesProto {
		const proto = pb.communicator_objects.UnityRLCapabilitiesProto.create({
			baseRLCapabilities: rlCaps.baseRLCapabilities,
			concatenatedPngObservations: rlCaps.concatenatedPngObservations,
			compressedChannelMapping: rlCaps.compressedChannelMapping,
			hybridActions: rlCaps.hybridActions,
			variableLengthObservation: rlCaps.variableLengthObservation,
			multiAgentGroups: rlCaps.multiAgentGroups,
		});
		return proto;
	}

	static capabilitiesFromProto(proto: pb.communicator_objects.UnityRLCapabilitiesProto): CCCRLCapabilities {
		return new CCCRLCapabilities(
			proto.baseRLCapabilities,
			proto.concatenatedPngObservations,
			proto.compressedChannelMapping,
			proto.hybridActions,
			proto.variableLengthObservation,
			proto.multiAgentGroups
		);
	}

	static toUnityRLInitParameters(
		inputProto: pb.communicator_objects.UnityRLInitializationInputProto
	): CCCRLInitParameters {
		return {
			seed: inputProto.seed,
			numAreas: inputProto.numAreas,
			pythonLibraryVersion: inputProto.packageVersion,
			pythonCommunicationVersion: inputProto.communicationVersion,
			trainerCapabilities: GrpcExtensions.capabilitiesFromProto(
				inputProto.capabilities as pb.communicator_objects.UnityRLCapabilitiesProto
			),
		};
	}

	static toTrainingEnvironmentInitializedEvent(
		inputProto: pb.communicator_objects.TrainingEnvironmentInitialized
	): TrainingEnvironmentInitializedEvent {
		return {
			trainerPythonVersion: inputProto.pythonVersion,
			trainingSessionGuid: inputProto.torchVersion,
			mLAgentsVersion: inputProto.mlagentsEnvsVersion,
			mLAgentsEnvsVersion: inputProto.mlagentsEnvsVersion,
			torchVersion: inputProto.torchVersion,
			torchDeviceType: inputProto.torchDeviceType,
			numEnvironments: inputProto.numEnvs,
			numEnvironmentParameters: inputProto.numEnvironmentParameters,
			runOptions: inputProto.runOptions,
		};
	}

	static toTrainingBehaviorInitializedEvent(
		inputProto: pb.communicator_objects.TrainingBehaviorInitialized
	): TrainingBehaviorInitializedEvent {
		let rewardSignals = RewardSignals.None;
		rewardSignals |= inputProto.extrinsicRewardEnabled ? RewardSignals.Extrinsic : RewardSignals.None;
		rewardSignals |= inputProto.gailRewardEnabled ? RewardSignals.Gail : RewardSignals.None;
		rewardSignals |= inputProto.curiosityRewardEnabled ? RewardSignals.Curiosity : RewardSignals.None;
		rewardSignals |= inputProto.rndRewardEnabled ? RewardSignals.Rnd : RewardSignals.None;

		let trainingFeatures = TrainingFeatures.None;
		trainingFeatures |= inputProto.behavioralCloningEnabled
			? TrainingFeatures.BehavioralCloning
			: TrainingFeatures.None;
		trainingFeatures |= inputProto.recurrentEnabled ? TrainingFeatures.Recurrent : TrainingFeatures.None;
		trainingFeatures |= inputProto.trainerThreaded ? TrainingFeatures.Threaded : TrainingFeatures.None;
		trainingFeatures |= inputProto.selfPlayEnabled ? TrainingFeatures.SelfPlay : TrainingFeatures.None;
		trainingFeatures |= inputProto.curriculumEnabled ? TrainingFeatures.Curriculum : TrainingFeatures.None;

		return {
			behaviorName: inputProto.behaviorName,
			trainerType: inputProto.trainerType,
			rewardSignalFlags: rewardSignals,
			trainingFeatureFlags: trainingFeatures,
			visualEncoder: inputProto.visualEncoder,
			numNetworkLayers: inputProto.numNetworkLayers,
			numNetworkHiddenUnits: inputProto.numNetworkHiddenUnits,
			config: inputProto.config,
		};
	}

	static toAgentActionList(proto: pb.communicator_objects.UnityRLInputProto.ListAgentActionProto): ActionBuffers[] {
		const agentActions: ActionBuffers[] = [];
		proto.value.forEach((value) => {
			const agentAction = this.toActionBuffers(value as pb.communicator_objects.AgentActionProto);
			agentActions.push(agentAction);
		});
		return agentActions;
	}

	static toActionBuffers(proto: pb.communicator_objects.AgentActionProto): ActionBuffers {
		return new ActionBuffers(proto.continuousActions, proto.discreteActions);
	}

	static toBrainParametersProto(
		actionSpec: ActionSpec,
		name: string,
		isTraining: boolean
	): pb.communicator_objects.BrainParametersProto {
		const brainParametersProto = new pb.communicator_objects.BrainParametersProto();
		brainParametersProto.brainName = name;
		brainParametersProto.actionSpec = GrpcExtensions.toActionSpecProto(actionSpec);

		const supportHybrid = Academy.instance.trainerCapabilities?.hybridActions;
		if (!supportHybrid) {
			actionSpec.checkAllContinuousOrDiscrete();
			if (actionSpec.numContinuousActions > 0) {
				brainParametersProto.vectorActionSizeDeprecated.push(actionSpec.numContinuousActions);
				brainParametersProto.vectorActionSpaceTypeDeprecated =
					pb.communicator_objects.SpaceTypeProto.continuous;
			} else if (actionSpec.numDiscreteActions > 0) {
				brainParametersProto.vectorActionSpaceTypeDeprecated = pb.communicator_objects.SpaceTypeProto.discrete;
				brainParametersProto.vectorActionSizeDeprecated = actionSpec.branchSizes;
			}
		}

		return brainParametersProto;
	}

	static fromActionSpec(actionSpec: ActionSpec): pb.communicator_objects.ActionSpecProto {
		const actionSpecProto = new pb.communicator_objects.ActionSpecProto();
		actionSpecProto.numContinuousActions = actionSpec.numContinuousActions;
		if (actionSpec.branchSizes && actionSpec.branchSizes.length > 0) {
			actionSpecProto.discreteBranchSizes = actionSpec.branchSizes;
		}
		return actionSpecProto;
	}

	static toActionSpecProto(actionSpec: ActionSpec): pb.communicator_objects.ActionSpecProto {
		const actionSpecProto = new pb.communicator_objects.ActionSpecProto();
		actionSpecProto.numContinuousActions = actionSpec.numContinuousActions;
		actionSpecProto.numDiscreteActions = actionSpec.numDiscreteActions;

		if (actionSpec.branchSizes.length > 0) {
			actionSpecProto.discreteBranchSizes = actionSpec.branchSizes;
		}

		return actionSpecProto;
	}

	public static DemonstrationMetaDataToProto(
		metaData: DemonstrationMetaData
	): pb.communicator_objects.DemonstrationMetaProto {
		const proto = new pb.communicator_objects.DemonstrationMetaProto();
		proto.numberEpisodes = metaData.numberEpisodes;
		proto.numberSteps = metaData.numberSteps;
		proto.meanReward = metaData.meanReward;
		proto.demonstrationName = metaData.demonstrationName;
		return proto;
	}

	public static toDemonstrationMetaData(
		proto: pb.communicator_objects.DemonstrationMetaProto
	): DemonstrationMetaData {
		const metaData = new DemonstrationMetaData();
		metaData.numberEpisodes = proto.numberEpisodes;
		metaData.numberSteps = proto.numberSteps;
		metaData.meanReward = proto.meanReward;
		metaData.demonstrationName = proto.demonstrationName;
		return metaData;
	}

	public static brainParametersToProto(
		brainParameters: BrainParameters,
		name: string,
		isTraining: boolean
	): pb.communicator_objects.BrainParametersProto {
		const proto = new pb.communicator_objects.BrainParametersProto();
		proto.vectorActionSpaceTypeDeprecated = brainParameters.vectorActionSpaceType as number;
		proto.brainName = name;
		proto.isTraining = isTraining;
		proto.actionSpec = GrpcExtensions.toActionSpecProto(brainParameters.actionSpec);

		if (brainParameters.vectorActionSize != null) {
			proto.vectorActionSizeDeprecated.push(...brainParameters.vectorActionSize);
		}
		if (brainParameters.vectorActionDescriptions != null) {
			proto.vectorActionDescriptionsDeprecated.push(...brainParameters.vectorActionDescriptions);
		}
		return proto;
	}

	public static toInfoActionPairProto(ai: AgentInfo): pb.communicator_objects.AgentInfoActionPairProto {
		const agentInfoProto = this.toAgentInfoProto(ai); // 假设这个方法已经返回pb.communicator_objects.AgentInfoProto类型
		const agentActionProto = new pb.communicator_objects.AgentActionProto();

		if (!ai.storedActions.isEmpty()) {
			if (!ai.storedActions.continuousActions.isEmpty()) {
				agentActionProto.continuousActions = ai.storedActions.continuousActions.data;
			}
			if (!ai.storedActions.discreteActions.isEmpty()) {
				agentActionProto.discreteActions = ai.storedActions.discreteActions.data;
			}
		}

		const proto = new pb.communicator_objects.AgentInfoActionPairProto();
		proto.agentInfo = agentInfoProto;
		proto.actionInfo = agentActionProto;
		return proto;
	}

	public static toAgentInfoProto(ai: AgentInfo): pb.communicator_objects.AgentInfoProto {
		if (ai.groupId > 0) {
			const trainerCanHandle = Academy.instance.trainerCapabilities?.multiAgentGroups;
			if (!trainerCanHandle) {
				if (this.hasWarnedAboutMultiAgentGroups) {
					this.hasWarnedAboutMultiAgentGroups = true;
				}
			}
		}

		const proto = new pb.communicator_objects.AgentInfoProto();
		proto.done = ai.done;
		proto.reward = ai.reward;
		proto.groupId = ai.groupId;
		proto.maxStepReached = ai.maxStepReached;
		proto.id = ai.episodeId;
		proto.groupId = ai.groupId; // 这行重复，可以考虑删除

		if (ai.discreteActionMasks != null) {
			proto.actionMask = [...ai.discreteActionMasks];
		}
		return proto;
	}

	private static _haveWarnedTrainerCapabilitiesMultiPng: boolean = false;
	private static _haveWarnedTrainerCapabilitiesMapping = false;

	static getObservationProto(
		sensor: ISensor,
		observationWriter: ObservationWriter
	): pb.communicator_objects.ObservationProto {
		const obsSpec = sensor.getObservationSpec();
		const shape = obsSpec.shape;
		let observationProto: pb.communicator_objects.ObservationProto | null = null;
		const compressionSpec = sensor.getCompressionSpec();
		let compressionType = compressionSpec.sensorCompressionType;

		if (compressionType === SensorCompressionType.PNG && shape.length === 3 && shape[2] > 3) {
			const trainerCanHandle =
				!Academy.instance.trainerCapabilities ||
				Academy.instance.trainerCapabilities.concatenatedPngObservations;
			if (!trainerCanHandle) {
				if (!this._haveWarnedTrainerCapabilitiesMultiPng) {
					console.warn(
						`Attached trainer doesn't support multiple PNGs. Switching to uncompressed observations for sensor ${sensor.getName()}. ` +
							'Please find the versions that work best together from our release page: ' +
							'https://github.com/Unity-Technologies/ml-agents/releases'
					);
					this._haveWarnedTrainerCapabilitiesMultiPng = true;
				}
				compressionType = SensorCompressionType.None;
			}
		}

		if (compressionType !== SensorCompressionType.None && shape.length === 3 && shape[2] > 3) {
			const trainerCanHandleMapping =
				!Academy.instance.trainerCapabilities || Academy.instance.trainerCapabilities.compressedChannelMapping;
			const isTrivialMapping = compressionSpec.isTrivialMapping();
			if (!trainerCanHandleMapping && !isTrivialMapping) {
				if (!this._haveWarnedTrainerCapabilitiesMapping) {
					console.warn(
						`The sensor ${sensor.getName()} is using non-trivial mapping and ` +
							"the attached trainer doesn't support compression mapping. " +
							'Switching to uncompressed observations. ' +
							'Please find the versions that work best together from our release page: ' +
							'https://github.com/Unity-Technologies/ml-agents/releases'
					);
					this._haveWarnedTrainerCapabilitiesMapping = true;
				}
				compressionType = SensorCompressionType.None;
			}
		}

		if (compressionType === SensorCompressionType.None) {
			const numFloats = SensorExtensions.observationSize(sensor);
			const floatDataProto = new pb.communicator_objects.ObservationProto.FloatData();
			for (let i = 0; i < numFloats; i++) {
				floatDataProto.data.push(0.0);
			}

			observationWriter.setTarget(floatDataProto.data, sensor.getObservationSpec(), 0);
			sensor.write(observationWriter);

			observationProto = new pb.communicator_objects.ObservationProto({
				floatData: floatDataProto,
				compressionType: pb.communicator_objects.CompressionTypeProto.NONE,
			});
		} else {
			const compressedObs = sensor.getCompressedObservation();
			if (compressedObs === null) {
				throw new Error(
					`GetCompressedObservation() returned null data for sensor named ${sensor.getName()}. ` +
						"You must return a byte array. If you don't want to use compressed observations, " +
						'return CompressionSpec.default() from getCompressionSpec().'
				);
			}
			observationProto = new pb.communicator_objects.ObservationProto({
				compressedData: new Uint8Array(compressedObs),
				compressionType: compressionSpec.sensorCompressionType as number,
			});
			if (compressionSpec.compressedChannelMapping !== null) {
				observationProto.compressedChannelMapping.push(
					...(compressionSpec.compressedChannelMapping as number[])
				);
			}
		}

		const dimensionProperties = obsSpec.dimensionProperties;
		for (let i = 0; i < dimensionProperties.length; i++) {
			observationProto.dimensionProperties.push(dimensionProperties[i] as number);
		}

		if (
			dimensionProperties[0] == DimensionProperty.VariableSize &&
			dimensionProperties[1] == DimensionProperty.None
		) {
			const trainerCanHandleVarLenObs =
				!Academy.instance.trainerCapabilities || Academy.instance.trainerCapabilities.variableLengthObservation;
			if (!trainerCanHandleVarLenObs) {
				throw new Error('Variable Length Observations are not supported by the trainer');
			}
		}

		for (let i = 0; i < shape.length; i++) {
			observationProto.shape.push(shape[i]);
		}

		const sensorName = sensor.getName();
		if (sensorName !== '') {
			observationProto.name = sensorName;
		}

		observationProto.observationType = obsSpec.observationType as number;
		return observationProto;
	}
}
